#!/usr/bin/env python
# -*- coding: utf-8 -*-

from base64 import b64encode, b64decode

import md5
import os, os.path, sys, time
import tempfile

from potherb import log
from potherb.cellstore import CellStore
from potherb.help import heappush, heappop

#######################################################################

def mergeUsingHeap(store_merge, store_cells) :
    cells = []
    for cell in store_cells :
        try :
            k, v = cell.next()
            heappush(cells, (k, v, cell))
        except :
            # empty store
            continue

    count_records_merged = 0
    while cells :
        _k, _v, cell_to_merge = heappop(cells)
        if not cell_to_merge or not _k or not _v : break

        store_merge.append(_k, _v)
        count_records_merged += 1

        try :
            k, v = cell_to_merge.next()
            heappush(cells, (k, v, cell_to_merge))
        except :
            #print sys.exc_info()
            pass
        # while

    log.info('Heap   Merge: %d records merged.', count_records_merged)
    # mergeUsingHeap()

def mergeUsingBubble(store_merge, store_cells) :
    cells = {}
    for store in store_cells :
        try :
            kv = store.next()
            cells[ store ] = kv
        except :
            # empty store
            continue

    count_records_merged = 0
    while cells :
        cell_to_merge = None
        kv_to_merge = ('','')       # merge keep descend order
        for _cell, _kv in cells.iteritems() :
            if _kv > kv_to_merge :
                kv_to_merge = _kv
                cell_to_merge = _cell
        if not cell_to_merge : break

        _k, _v = kv_to_merge
        store_merge.append(_k, _v)
        count_records_merged += 1

        try :
            kv = cell_to_merge.next()
            cells[ cell_to_merge ] = kv
        except :
            del cells[ cell_to_merge ]

        # while

    log.info('Bubble Merge: %d records merged.', count_records_merged)
    # mergeUsingBubble()

def test1(cell_size = 65536) :
    records_total = 1048576
    records_cell  = cell_size

    tmpstore_files = []
    tmpmerge_files = []

    time_prepare = 0
    time_sort  = 0
    time_save  = 0
    time_merge = 0
    time_total = 0

    min_key = 1
    max_key = min_key + records_total
    k = min_key
    while k < max_key :
        # STEP. prepare
        time_start = time.time()
        haystack = []
        count_cell = 0
        while count_cell < records_cell and k < max_key :
            data = '%16d' % k
            key = md5.new(data).hexdigest()
            haystack.append((key, data))

            k += 1
            count_cell += 1
        time_stop = time.time()
        time_last = time_stop - time_start
        time_prepare += time_last

        # STEP. sorted
        time_start = time.time()
        haystack.sort(reverse=True)
        time_stop  = time.time()
        time_last  = time_stop - time_start
        time_sort += time_last

        # STEP. save
        time_start = time.time()

        tempfd, tempname = tempfile.mkstemp(suffix='.store', prefix='cell_')
        os.close(tempfd)
        log.debug('CELL  temporary file : %s', tempname)
        tmpstore_files.append( tempname )
        store = CellStore(tempname)
        for _k,_v in haystack :
            store.append(_k,_v)
        del store

        time_stop  = time.time()
        time_last  = time_stop - time_start
        time_save += time_last

        # while loop

    # STEP. merge
    for merge,merge_name in ((mergeUsingHeap, 'HEAP'), (mergeUsingBubble, 'BUBBLE')) :
        time_start = time.time()

        cells = []
        for cellfile in tmpstore_files :
            store = CellStore(cellfile)
            cells.append(store)

        tempfd, tempname = tempfile.mkstemp(suffix='.store', prefix='merge_')
        os.close(tempfd)
        log.debug('MERGE temporary file : %s', tempname)
        tmpmerge_files.append( tempname )
        store_merge = CellStore(tempname)

        merge(store_merge, cells)

        del store_merge
        for store in cells :
            del store
        cells = []

        time_stop  = time.time()
        time_last  = time_stop - time_start
        time_merge += time_last
        log.debug('Merge (use %s) last %.2f seconds', merge_name, time_last)

    for tempname in tmpstore_files :
        os.remove(tempname)
    for tempname in tmpmerge_files :
        os.remove(tempname)

    time_total = time_prepare + time_sort + time_save + time_merge

    tps_prepare = records_total / (time_prepare if time_prepare else 0.001)
    tps_sort    = records_total / (time_sort    if time_sort    else 0.001)
    tps_save    = records_total / (time_save    if time_save    else 0.001)
    tps_merge   = records_total / (time_merge   if time_merge   else 0.001)
    tps         = records_total / (time_total   if time_total   else 0.001)

    log.info('Total count of records: %d', records_total)
    log.info('Records count per CELL: %d', records_cell)
    log.info('Time-Last (in seconds) prepare=%.2f sort=%.2f save=%.2f merge=%.2f TOTAL=%.2f', time_prepare, time_sort, time_save, time_merge, time_total)
    log.info('TPS prepare=%.2f sort=%.2f save=%.2f merge=%.2f avg=%.2f', tps_prepare, tps_sort, tps_save, tps_merge, tps)
    # test1()

#######################################################################

if  __name__ == '__main__' :
    log.enableDebug()

    test1(65536)
    test1(131072)
    test1(262144)

